Hello!

Pull Request overview

• Allow loading custom modules

Details

This is an experimental branch to allow models with custom architectures to integrate with Sentence Transformers. Behind the scenes, Sentence Transformer models always rely on a list of modules, with configuration in modules.json. Usually, this contains a Transformer followed by a Pooler and optionally a Normalize and/or Dense module, but we can now fully change this up.

I have two implementation formats that should both work:

1. tomaarsen/jina-clip-v1-st hosts all custom modeling files, notably custom_st.py. This class is loaded via the modules.json. This class handles tokenization/preprocessing of texts/images, and implements the forward. The forward method accepts a dictionary of features from the tokenization, and then must output a dictionary. Some modules might output sentence_embedding, while others may output token_embeddings. The latter can then still be used with the Pooling module, which will convert the token embeddings into sentence embeddings.
2. tomaarsen/jina-clip-v1-st-remote and tomaarsen/jina-clip-implementation-st are similar, but here the modeling file is in a separate repository (the latter). The modules.json file now uses tomaarsen/jina-clip-implementation-st--custom_st.Transformer, and this repository hosts custom_st.py. Otherwise, everything is the same as the other implementation.

Note: you don't have to use Transformer or custom_st as names, you can be flexible here.

Usage

You can now use these custom modules in Sentence Transformers directly via trust_remote_code=True:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("tomaarsen/jina-clip-v1-st-remote", trust_remote_code=True)
# or
# model = SentenceTransformer("tomaarsen/jina-clip-v1-st", trust_remote_code=True)

# New meaningful sentences
sentences = ['A blue cat', 'A red cat']

# Public image URLs
image_urls = [
    'https://i.pinimg.com/600x315/21/48/7e/21487e8e0970dd366dafaed6ab25d8d8.jpg',
    'https://i.pinimg.com/736x/c9/f2/3e/c9f23e212529f13f19bad5602d84b78b.jpg'
]

text_embeddings = model.encode(sentences)
image_embeddings = model.encode(image_urls)

# Compute similarities
print(model.similarity(text_embeddings[0], text_embeddings[1])) # text embedding similarity
print(model.similarity(text_embeddings, image_embeddings)) # text-image cross-modal similarity
"""
tensor([[0.5636]])
tensor([[0.2906, 0.0569],
        [0.1277, 0.2916]]
"""

I believe this opens up a lot of opportunities for finetuning custom architectures, too. Even modalities not previously used with Sentence Transformers (e.g. audio?) should be totally feasible.

cc @bwanglzu

• Tom Aarsen